#include "sat_cl.h"
#include "util.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>

void setup_cl(SAT_test_t *t);
void setup_buffers(SAT_test_t * t, cl_float4(*gen)(int i));
void warmup(SAT_test_t * t, int num_iters);
bool run(SAT_test_t * t, int num_iters);
bool check(SAT_test_t *t);
	void SAT_gold(SAT_test_t *t);
	void print_arr(SAT_test_t * t, cl_float4 *arr);
void post(SAT_test_t *t, int run_num_iters, bool result);
void cleanup(SAT_test_t *t);

#define PTX_OUTPUT

/////////////////////////////////////////////////////

void setup_buffers(SAT_test_t * t, cl_float4(*gen)(int i))
{
	// malloc host memory
	t->hst_src = (cl_float4 *)malloc(t->num * t->num * sizeof(cl_float4));
	t->hst_dst = (cl_float4 *)malloc(t->num * t->num * sizeof(cl_float4));
	for (int i = 0; i < t->num * t->num; i++) {
		t->hst_src[i] = gen(i);
	}

	cl_int errcode_ret;
	// create GPU buffers
	printf("Creating GPU buffers...\n");
	t->dev_src = clCreateBuffer(t->hContext, CL_MEM_READ_WRITE, t->num * t->num * sizeof(cl_float4), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer dev_src");
	t->dev_dst = clCreateBuffer(t->hContext, CL_MEM_READ_WRITE, t->num * t->num * sizeof(cl_float4), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer dev_dst");
}

void write_ptx(cl_program program, char *filename)
{
	cl_int errcode_ret;
	size_t progSize;
	errcode_ret = clGetProgramInfo(program, CL_PROGRAM_BINARY_SIZES, sizeof(size_t), &progSize, NULL);
	ERR_CHECK(errcode_ret, "clGetProgramInfo");
	char ** prog = (char **) malloc (sizeof(char **));
	prog[0] = (char *) malloc(progSize * sizeof(char *));
	errcode_ret = clGetProgramInfo(program, CL_PROGRAM_BINARIES, sizeof(char **) * progSize, prog, NULL);
	ERR_CHECK(errcode_ret, "clGetProgramInfo");
	FILE * f = fopen(filename, "w");
	fprintf(f, "%s\n", prog[0]);
	fclose(f);
}

void setup_cl(SAT_test_t *t)
{
	cl_int errcode_ret;

	// Get OpenCL platform count
	cl_uint NumPlatforms;
	clGetPlatformIDs (0, NULL, &NumPlatforms);

	// Get all OpenCL platform IDs
	cl_platform_id* PlatformIDs;
	PlatformIDs = new cl_platform_id[NumPlatforms];
	clGetPlatformIDs (NumPlatforms, PlatformIDs, NULL);

	// find NVIDIA & AMD platforms
	char cBuffer[1024];
	cl_int NvPlatform = -1;
	cl_int AMDPlatform = -1;
	for(cl_uint i = 0; i < NumPlatforms; ++i)
	{
		clGetPlatformInfo (PlatformIDs[i], CL_PLATFORM_NAME, 1024, cBuffer, NULL);
		printf("%s\n", cBuffer);
		if(strstr(cBuffer, "NVIDIA") != NULL) {
			NvPlatform = i;
		}
		else if (strstr(cBuffer, "AMD") != NULL) {
			AMDPlatform = i;
		}
	}

	// check for AMD and NVIDIA GPU devices
	cl_device_id cdDevice;
	cl_uint NvNumDevices = 0;
	cl_uint AMDNumDevices = 0;
	if (AMDPlatform != -1)
		clGetDeviceIDs(PlatformIDs[AMDPlatform], CL_DEVICE_TYPE_GPU, 0, NULL, &AMDNumDevices);
	if (NvPlatform != -1)
		clGetDeviceIDs(PlatformIDs[NvPlatform], CL_DEVICE_TYPE_GPU, 0, NULL, &NvNumDevices);

	// if there is an AMD GPU, take it, or take an NVIDIA GPU if it is there
	if (AMDNumDevices > 0)
		clGetDeviceIDs(PlatformIDs[AMDPlatform], CL_DEVICE_TYPE_GPU, 1, &cdDevice, NULL);
	else if (NvNumDevices > 0)
		clGetDeviceIDs(PlatformIDs[NvPlatform], CL_DEVICE_TYPE_GPU, 1, &cdDevice, NULL);
	else {
		fprintf(stderr, "could not find any GPU devices, exiting\n");
		delete [] PlatformIDs;
		exit(-1);
	}

	delete [] PlatformIDs;

	clGetDeviceInfo(cdDevice, CL_DEVICE_MAX_WORK_GROUP_SIZE, sizeof(size_t), &(t->maxBlockSize), NULL);
	//Create a context
	printf("Creating context for %s GPU...\n", (AMDNumDevices > 0) ? "AMD" : "NVIDIA");
	t->hContext = clCreateContext(NULL, 1, &cdDevice, NULL, NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateContext");

	size_t nContextDescriptorSize; 
	clGetContextInfo(t->hContext, CL_CONTEXT_DEVICES, 0, 0, &nContextDescriptorSize); 
	cl_device_id * aDevices = (cl_device_id *) malloc(nContextDescriptorSize); 
	clGetContextInfo(t->hContext, CL_CONTEXT_DEVICES, nContextDescriptorSize, aDevices, 0);

	// create a command queue for first
	// device the context reported  
	t->hCmdQueue = clCreateCommandQueue(t->hContext, aDevices[0], CL_QUEUE_PROFILING_ENABLE, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateCommandQueue");
	free(aDevices);
	
	// create & compile SAT program
	cl_program hSATScanProg; 
	char * SATScanCode;
	size_t SATScanLen;
	
	printf("Compiling scan program...\n");
	SATScanCode = oclLoadProgSource("scan_kernel.cl", "", &SATScanLen);
	hSATScanProg = clCreateProgramWithSource(t->hContext,1, (const char **)&SATScanCode, &SATScanLen, &errcode_ret); 
	ERR_CHECK(errcode_ret, "clCreateProgramWithSource SATScan");
	errcode_ret = clBuildProgram(hSATScanProg, 0, NULL, NULL, NULL, NULL);
	ERR_CHECK(errcode_ret, "clBuildProgram SATScan");
	free(SATScanCode);

#ifdef PTX_OUTPUT
	write_ptx(hSATScanProg, "scan_cl.ptx");
#endif

	// create SAT kernel 
	printf("Creating scan kernel...\n");
	t->hSATScanKernel = clCreateKernel(hSATScanProg, "kernel_sat_scan", &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateKernel SATScan");

	// create & compile transpose program
	cl_program hTransposeProg; 
	char * TransposeCode;
	size_t TransposeLen;
	
	printf("Compiling transpose program...\n");
	TransposeCode = oclLoadProgSource("transpose_kernel.cl", "", &TransposeLen);
	hTransposeProg = clCreateProgramWithSource(t->hContext,1, (const char **)&TransposeCode, &TransposeLen, &errcode_ret); 
	ERR_CHECK(errcode_ret, "clCreateProgramWithSource Transpose");
	errcode_ret = clBuildProgram(hTransposeProg, 0, NULL, NULL, NULL, NULL);
	ERR_CHECK(errcode_ret, "clBuildProgram Transpose");
	free(TransposeCode);

#ifdef PTX_OUTPUT
	write_ptx(hTransposeProg, "transpose_cl.ptx");
#endif

	// create kernel 
	printf("Creating transpose kernel...\n");
	t->hTransposeKernel = clCreateKernel(hTransposeProg, "kernel_naive_transpose", &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateKernel Transpose");
} 

cl_float4 gen_float4(float i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = i;
	return a;
}

cl_float4 gen0(int i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = 1.0f;
	return a;
}

cl_float4 gen1(int i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = 1.0f * i;
	return a;
}

cl_float4 gen2(int i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = 2.0f * i;
	return a;
}

cl_float4 gen3(int i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = 0.0f;
	return a;
}

cl_float4 gen4(int i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = cos(i/1000.0f);
	return a;
}

cl_float4 gen5(int i) {
	cl_float4 a;
	a.s[0] = a.s[1] = a.s[2] = a.s[3] = sin(i/100.0f);
	return a;
}

void warmup(SAT_test_t * t, int num_iters)
{
	for (int i = 0; i < num_iters; i++) {
		launch_SAT(t);
		printf("Warmup %d/%d complete\n", i+1, num_iters);
	}
}

bool run(SAT_test_t * t, int num_iters)
{
	t->time = t->transpose_time = t->scan_time = 0.0;

	for (int i = 0; i < num_iters; i++) {
		double time = launch_SAT(t);
		t->time += time;
		if (!check(t)) 
			return false;
		printf("Run %d complete\n", i+1, num_iters);
	}
	return true;
}

bool check(SAT_test_t *t)
{
	float EPSILON = 20.0f;
	cl_int errcode_ret;

	SAT_gold(t);
	cl_float4 * temp = (cl_float4 *) malloc (t->num * t->num * sizeof(cl_float4));
	errcode_ret = clEnqueueReadBuffer(t->hCmdQueue, t->dev_dst, CL_TRUE, 0, t->num * t->num * sizeof(cl_float4),
		temp, 0, NULL, NULL);
	ERR_CHECK(errcode_ret, "clEnqueueReadBuffer");
	//print_arr(t, temp);
	//print_arr(t, t->hst_dst);
	for (int i = 0; i < t->num * t->num; i++) {
		cl_float4 hst = t->hst_dst[i];
		cl_float4 dev = temp[i];
		if (abs(hst.s[0] - dev.s[0]) >= EPSILON ||
			abs(hst.s[1] - dev.s[1]) >= EPSILON ||
			abs(hst.s[2] - dev.s[2]) >= EPSILON ||
			abs(hst.s[3] - dev.s[3]) >= EPSILON ) {
				fprintf(stderr, "check failed at %d: [%f, %f, %f, %f] != [%f, %f, %f, %f]\n", 
					i,
					hst.s[0], hst.s[1], hst.s[2], hst.s[3],
					dev.s[0], dev.s[1], dev.s[2], dev.s[3]);
				free(temp);
				return false;
		}
	}
	free(temp);
	return true;
}

void SAT_gold(SAT_test_t * t)
{
	for (int y = 0; y < t->num; y++) {
		for (int x = 0; x < t->num; x++) {
			int i = y * t->num + x;
			cl_float4 me = t->hst_src[i];
			cl_float4 below = (y < t->num-1) ? t->hst_dst[(y+1)*t->num + x] : gen_float4(0.0f);
			cl_float4 above = (y > 0) ? t->hst_dst[(y-1)*t->num + x] : gen_float4(0.0f);
			cl_float4 left = (x > 0) ? t->hst_dst[y*t->num + x-1] : gen_float4(0.0f);
			cl_float4 below_left = (y < t->num-1 && x > 0) ? 
				t->hst_dst[(y+1)*t->num + x-1] : gen_float4(0.0f);
			cl_float4 above_left = (y > 0 && x > 0) ? 
				t->hst_dst[(y-1)*t->num + x-1] : gen_float4(0.0f);

			t->hst_dst[i].s[0] = me.s[0] + above.s[0] + left.s[0] - above_left.s[0];
			t->hst_dst[i].s[1] = me.s[1] + above.s[1] + left.s[1] - above_left.s[1];
			t->hst_dst[i].s[2] = me.s[2] + above.s[2] + left.s[2] - above_left.s[2];
			t->hst_dst[i].s[3] = me.s[3] + above.s[3] + left.s[3] - above_left.s[3];
			
		}
	}
}

void print_arr(SAT_test_t * t, cl_float4 *arr)
{
	for (int i = 0; i < t->num * t->num; i++) {
		if (i % t->num == 0)
			printf("[");
		printf("%3.0f", arr[i].s[0]);
		if (i % t->num < t->num-1)
			printf(", ");
		else
			printf("]\n");
	}
}

void post(SAT_test_t * t, int run_num_iters, bool result)
{
	printf("----------------RESULTS----------------\n");
	printf("PASS? %s\n", result ? "YES" : "NO");
	printf("Size: %d x %d\n", t->num, t->num);
	printf("Time: %f\n", t->time / run_num_iters);
}

void output(SAT_test_t *t, FILE *csv, int run_num_iters) {
	fprintf(csv, "%u, %u, %Lf, %Lf, %Lf\n",
		t->num,
		t->blockSize,
		t->scan_time/(2 * run_num_iters),
		t->transpose_time/(2 * run_num_iters),
		t->time / run_num_iters);
	// / 2 since scan and transpose done twice
}

void cleanup(SAT_test_t * t)
{
	cl_int errcode_ret;
	free(t->hst_src);
	free(t->hst_dst);
	clFinish(t->hCmdQueue);
	errcode_ret = clReleaseMemObject(t->dev_src);
	ERR_CHECK(errcode_ret, "clReleaseMemObject dev_src");
	errcode_ret = clReleaseMemObject(t->dev_dst);
	ERR_CHECK(errcode_ret, "clReleaseMemObject dev_dst");
}

void test_SAT(int sizeMin, int sizeMax,
			  int warmup_num_iters, int run_num_iters, 
			  size_t blockMin, size_t blockMax, 
			  bool split, char *filename, cl_float4(*gen)(int i))
{
	SAT_test_t t;
	t.num = sizeMin;
	t.blockSize = blockMin;
	t.split = split;
	bool result = false;

	FILE * csv = fopen(filename, "a");
	if (csv == NULL)
		fprintf(stderr, "can't open output file\n");

	if (csv != NULL)
		fprintf(csv, "size, block size, scan time, transpose time, total time\n");

	setup_cl(&t);
	setup_buffers(&t, gen);

	if (t.maxBlockSize < blockMin)
		blockMin = t.maxBlockSize;
	if (t.maxBlockSize < blockMax)
		blockMax = t.maxBlockSize;

	printf("blockMIN: %d blockMAX: %d t.maxBlockSize: %d\n", blockMin, blockMax, t.maxBlockSize);
	warmup(&t, warmup_num_iters);

	for (t.num = sizeMin; t.num <= sizeMax; t.num <<= 1) {
		for (t.blockSize = blockMin; t.blockSize <= blockMax; t.blockSize += 128) {
			setup_buffers(&t, gen);
			result = run(&t, run_num_iters);
			post(&t, run_num_iters, result);
			if (result && csv != NULL)
				output(&t, csv, run_num_iters);
			cleanup(&t);
		}
	}

	if (csv != NULL)
		fclose(csv);
}

int main(int argc, char *argv[])
{
	char *filename = "output.csv";
	test_SAT(256, 4096, 1, 5, 128, 256, true, filename, gen4);
	test_SAT(256, 4096, 1, 5, 128, 256, false, filename, gen4);

	return 0;
}
